# Copyright (c) 2017 Presto Labs Pte. Ltd.
# Author: inkyu

import ujson as json

from coin.base.book.multiple_book_builder_wrapper \
    import MultipleBookBuilderWrapper
from coin.exchange.bitmex.book_builder import BookBuilder
from coin.exchange.bitmex.kr_rest.futures_product import BitmexFuturesProduct
from coin.exchange.base.ws.handler import WebsocketHandler


class BitmexMultipleBookHandler(WebsocketHandler):
  def __init__(self, aggregated_ticksizes=None):
    super().__init__()
    self._instruments = {}
    self._aggregated_ticksizes = aggregated_ticksizes

  def on_ws_orderbook_l2(self, record, queue_data, topic_data):
    message = record.data
    if not isinstance(message, dict):
      message = json.loads(message.tobytes())

    updated_books = {}
    action = message['action']
    if 'table' in message and message['table'] == 'orderBookL2':
      if action == 'partial':
        assert message['keys'] == ['symbol', 'id', 'side']
        self._books = {}
        for data in message['data']:
          symbol = data['symbol']
          product = BitmexFuturesProduct.FromStrNativeProductNothrow(symbol)
          if product is None:
            continue
          if not self.is_subscribed('book', symbol):
            continue
          if symbol not in self._books:
            self._books[symbol] = MultipleBookBuilderWrapper(BookBuilder(symbol),
                                                             self._aggregated_ticksizes)
          book = self._books[symbol]
          diff = book.get_original_book().get_insert_diff(data)
          book.update_aggregated_books(diff[0], diff[1], diff[2])
          book.get_original_book().partial(data)
          book.update_timestamp(record.timestamp)
          updated_books[symbol] = book
        self.ready = True

      elif action == 'delete':
        for data in message['data']:
          symbol = data['symbol']
          if symbol in self._books:
            book = self._books[symbol]
            diff = book.get_original_book().get_update_diff(data['id'], data['side'], 0)
            if diff:
              book.update_aggregated_books(diff[0], diff[1], diff[2])
            book.get_original_book().delete(data['id'], data['side'])
            book.update_timestamp(record.timestamp)
            updated_books[symbol] = book

      elif action == 'insert':
        for data in message['data']:
          symbol = data['symbol']
          if symbol in self._books:
            book = self._books[symbol]
            diff = book.get_original_book().get_insert_diff(data)
            book.update_aggregated_books(diff[0], diff[1], diff[2])
            book.get_original_book().insert(data)
            book.update_timestamp(record.timestamp)
            updated_books[symbol] = book

      elif action == 'update':
        for data in message['data']:
          symbol = data['symbol']
          if symbol in self._books:
            book = self._books[symbol]
            diff = book.get_original_book().get_update_diff(data['id'], data['side'], data['size'])
            if diff:
              book.update_aggregated_books(diff[0], diff[1], diff[2])
            book.get_original_book().update(data['id'], data['side'], data['size'])
            book.update_timestamp(record.timestamp)
            updated_books[symbol] = book

      else:
        raise ValueError('Unknown action: %s' % action)

    for book in updated_books.values():
      try:
        if book.bid0().price < book.ask0().price:
          self.publish_event('book', book.symbol, book)
      except IndexError:
        continue
